import argparse
import logging
import os
import shutil
from datetime import datetime, timedelta
from typing import List

import torch
from omegaconf import OmegaConf
from torch.utils.data import ConcatDataset, DataLoader, Subset
from tqdm import tqdm

from marigold_inpaint.controlnet import ControlNetMarigoldPipeline
from src.dataset import BaseDepthDataset, DatasetMode, get_dataset
from src.dataset.mixed_sampler import MixedBatchSampler
from marigold_inpaint.controlnet_marigold_trainer import ControlNetMarigoldTrainer
from src.util.config_util import (
    find_value_in_omegaconf,
    recursive_load_config,
)
from src.util.depth_transform import (
    DepthNormalizerBase,
    get_depth_normalizer,
)
from src.util.logging_util import (
    config_logging,
    init_wandb,
    load_wandb_job_id,
    log_slurm_job_id,
    save_wandb_job_id,
    tb_logger,
)
from marigold_inpaint.base_dataset import BaseDataset
from src.util.slurm_util import get_local_scratch_dir, is_on_slurm
import sys
from accelerate import Accelerator
# # 模拟命令行参数
# sys.argv = [
#     "train_controlnet.py",  # 模拟脚本名称
#     "--config","config/train_controlnet_marigold.yaml",  # 模拟配置文件路径
    
# ]

if "__main__" == __name__:
    t_start = datetime.now()
    print(f"start at {t_start}")

    # -------------------- Arguments --------------------
    parser = argparse.ArgumentParser(description="Train your cute model!")
    parser.add_argument(
        "--config",
        type=str,
        default="config/train_controlnet_marigold.yaml",
        help="Path to config file.",
    )
    parser.add_argument(
        "--resume_run",
        action="store",
        default=None,
        help="Path of checkpoint to be resumed. If given, will ignore --config, and checkpoint in the config",
    )
    parser.add_argument(
        "--output_dir", type=str, default=None, help="directory to save checkpoints"
    )
    parser.add_argument("--no_cuda", action="store_true", help="Do not use cuda.")
    parser.add_argument(
        "--exit_after",
        type=int,
        default=-1,
        help="Save checkpoint and exit after X minutes.",
    )
    parser.add_argument("--no_wandb", action="store_true", help="run without wandb")
    parser.add_argument(
        "--do_not_copy_data",
        action="store_true",
        help="On Slurm cluster, do not copy data to local scratch",
    )
    parser.add_argument(
        "--base_data_dir", type=str, default=None, help="directory of training data"
    )
    parser.add_argument(
        "--base_ckpt_dir",
        type=str,
        default=None,
        help="directory of pretrained checkpoint",
    )
    parser.add_argument(
        "--add_datetime_prefix",
        action="store_true",
        help="Add datetime to the output folder name",
    )

    args = parser.parse_args()
    resume_run = args.resume_run
    output_dir = args.output_dir

    # -------------------- Initialization --------------------
    # Resume previous run
    if resume_run is not None:
        print(f"Resume run: {resume_run}")
        out_dir_run = os.path.dirname(os.path.dirname(resume_run))
        job_name = os.path.basename(out_dir_run)
        # Resume config file
        #TODO: use saved config?
        # cfg = OmegaConf.load(os.path.join(out_dir_run, "config.yaml"))
        cfg = recursive_load_config(args.config)
    else:
        # Run from start
        cfg = recursive_load_config(args.config)
        # Full job name
        pure_job_name = os.path.basename(args.config).split(".")[0]
        # Add time prefix
        if args.add_datetime_prefix:
            job_name = f"{t_start.strftime('%y_%m_%d-%H_%M_%S')}-{pure_job_name}"
        else:
            job_name = pure_job_name

        # Output dir
        if output_dir is not None:
            out_dir_run = os.path.join(output_dir, job_name)
        else:
            out_dir_run = os.path.join("./output", job_name)
        os.makedirs(out_dir_run, exist_ok=True)
    # Other directories
    out_dir_ckpt = os.path.join(out_dir_run, "checkpoint")
    if not os.path.exists(out_dir_ckpt):
        os.makedirs(out_dir_ckpt, exist_ok=True)
    out_dir_tb = os.path.join(out_dir_run, "tensorboard")
    if not os.path.exists(out_dir_tb):
        os.makedirs(out_dir_tb, exist_ok=True)
    out_dir_eval = os.path.join(out_dir_run, "evaluation")
    if not os.path.exists(out_dir_eval):
        os.makedirs(out_dir_eval, exist_ok=True)
    out_dir_vis = os.path.join(out_dir_run, "visualization")
    if not os.path.exists(out_dir_vis):
        os.makedirs(out_dir_vis, exist_ok=True)

    # -------------------- Logging settings --------------------
    config_logging(cfg.logging, out_dir=out_dir_run)
    logging.debug(f"config: {cfg}")

    # -------------------- Gradient accumulation steps --------------------
    eff_bs = cfg.dataloader.effective_batch_size
    accumulation_steps = eff_bs / cfg.dataloader.max_train_batch_size
    assert int(accumulation_steps) == accumulation_steps
    accumulation_steps = int(accumulation_steps)

    logging.info(
        f"Effective batch size: {eff_bs}, accumulation steps: {accumulation_steps}"
    )

    # -------------------- Data --------------------
    loader_seed = cfg.dataloader.seed
    if loader_seed is None:
        loader_generator = None
    else:
        loader_generator = torch.Generator().manual_seed(loader_seed)

    # Training dataset
    depth_transform: DepthNormalizerBase = get_depth_normalizer(
        cfg_normalizer=cfg.depth_normalization
    )
    logging.debug("Augmentation: ", cfg.augmentation)
    train_dataset = BaseDataset(
        filename_ls_path="/mnt/new/liufenglin/lishiyang/Marigold_Inpaint/train.csv",
        dataset_dir="/mnt/new/liufenglin/DL3DV-ALL-960P",
        H=336,
        W=512,
        random_mask_prob=0.5,
    )
    val_dataset = BaseDataset(
        filename_ls_path="/mnt/new/liufenglin/lishiyang/Marigold_Inpaint/val_others.csv",
        dataset_dir="/mnt/new/liufenglin/DL3DV-ALL-960P",
        H=336,
        W=512,
        random_mask_prob=0.5,
    )
    sub_val_dataset = Subset(val_dataset, list(range(30))) 
    sub_vis_dataset = Subset(val_dataset, list(range(15))) 
    train_loader = DataLoader(
        dataset=train_dataset,
        batch_size=cfg.dataloader.max_train_batch_size,
        num_workers=cfg.dataloader.num_workers,
        shuffle=True,
        generator=loader_generator,
    )
    val_loader = DataLoader(
        dataset=sub_val_dataset,
        batch_size=1,
        num_workers=1,
        shuffle=True,
    )
    vis_loader = DataLoader(
        dataset=sub_vis_dataset,
        batch_size=1,
        num_workers=1,
        shuffle=True,
    )
    # -------------------- Model --------------------
    _pipeline_kwargs = cfg.pipeline.kwargs if cfg.pipeline.kwargs is not None else {}
    #TODO: change to cfg
    model = ControlNetMarigoldPipeline.from_pretrained(
        "checkpoint/marigold-inpaint",
        # os.path.join(base_ckpt_dir, cfg.model.pretrained_path),
        **_pipeline_kwargs
    )

    # -------------------- Trainer --------------------
    # Exit time
    if args.exit_after > 0:
        t_end = t_start + timedelta(minutes=args.exit_after)
        logging.info(f"Will exit at {t_end}")
    else:
        t_end = None
    
    
    accelerator: Accelerator = Accelerator(
        gradient_accumulation_steps=accumulation_steps,
    )
    if accelerator.is_main_process:
        tb_logger.set_dir(out_dir_tb)
    trainer = ControlNetMarigoldTrainer(
        cfg=cfg,
        model=model,
        train_dataloader=train_loader,
        base_ckpt_dir='checkpoint',
        out_dir_ckpt=out_dir_ckpt,
        out_dir_eval=out_dir_eval,
        out_dir_vis=out_dir_vis,
        tb_log_dir=out_dir_tb,
        val_dataloader=val_loader,
        vis_dataloader=vis_loader,
        accelerator=accelerator,
        resume_run = resume_run,
    )

    # -------------------- Training & Evaluation Loop --------------------
    try:
        trainer.train(t_end=t_end)
    except Exception as e:
        logging.exception(e)
